#include "hip/hip_runtime.h"
#include <hip/hip_runtime.h>
#include <iostream>

#define CHECK_CUDA_ERROR(call)                                                 \
  {                                                                            \
    const hipError_t error = call;                                            \
    if (error != hipSuccess) {                                                \
        std::cerr                                                              \
          << "Error: " << hipGetErrorString(error) << std::endl;              \
      exit(1);                                                                 \
    }                                                                          \
}

__global__ void vectorAdd(int n, float *x, float *y, float *z) {
  int index = threadIdx.x + blockIdx.x * blockDim.x;
  if (index < n) {
    z[index] = x[index] + y[index];
  }
}

int main() {
  int numDevices;
  CHECK_CUDA_ERROR(hipGetDeviceCount(&numDevices));
  std::cout << "Number of devices: " << numDevices << std::endl;

  if (numDevices == 0) {
    std::cerr << "No CUDA devices found." << std::endl;
    return 1;
  }

  for (int dev = 0; dev < numDevices; ++dev) {
    hipSetDevice(dev);

    int n = 50000;
    size_t size = n * sizeof(float);

    float *h_A = new float[n];
    float *h_B = new float[n];
    float *d_A;
    float *d_B;
    float *d_C;

    // Initialize data on the host
    for (int i = 0; i < n; ++i) {
      h_A[i] = sinf((float)i);
      h_B[i] = cosf((float)i);
    }

    // Allocate GPU buffers for three vectors (two input, one output).
    CHECK_CUDA_ERROR(hipMalloc((void **)&d_A, size));
    CHECK_CUDA_ERROR(hipMalloc((void **)&d_B, size));
    CHECK_CUDA_ERROR(hipMalloc((void **)&d_C, size));

    // Copy inputs to device
    CHECK_CUDA_ERROR(hipMemcpy(d_A, h_A, size, hipMemcpyHostToDevice));
    CHECK_CUDA_ERROR(hipMemcpy(d_B, h_B, size, hipMemcpyHostToDevice));

    // Launch a kernel on the GPU with one thread for each element.
    int blockSize = 256;
    int numBlocks = (n + blockSize - 1) / blockSize;
    vectorAdd<<<numBlocks, blockSize>>>(n, d_A, d_B, d_C);

    // Check for any errors launching the kernel
    CHECK_CUDA_ERROR(hipGetLastError());

    // hipDeviceSynchronize waits for the kernel to finish, and returns
    // any errors encountered during the launch.
    CHECK_CUDA_ERROR(hipDeviceSynchronize());

    // Copy output to host
    float *h_C = new float[n];
    CHECK_CUDA_ERROR(hipMemcpy(h_C, d_C, size, hipMemcpyDeviceToHost));

    // Verify results
    // (simple check, only first 5 elements and last 5 elements are printed)
    for (int i = 0; i < 5; ++i) {
      std::cout << "Result " << i << " = " << h_C[i] << std::endl;
    }
    for (int i = n - 5; i < n; ++i) {
      std::cout << "Result " << i << " = " << h_C[i] << std::endl;
    }

    // Cleanup
    delete[] h_A;
    delete[] h_B;
    delete[] h_C;
    CHECK_CUDA_ERROR(hipFree(d_A));
    CHECK_CUDA_ERROR(hipFree(d_B));
    CHECK_CUDA_ERROR(hipFree(d_C));
  }

  return 0;
}
